package com.init.system.filter;

import java.util.concurrent.TimeUnit;

public class TokenBucket {
    private final long maxTokens; // 桶的最大令牌数
    private final long refillTokens; // 每次填充的令牌数
    private final long refillInterval; // 填充的时间间隔（毫秒）
    
    private long availableTokens; // 当前可用的令牌数
    private long lastRefillTime; // 上次填充时间

    public TokenBucket(long maxTokens, long refillTokens, long refillInterval) {
        this.maxTokens = maxTokens;
        this.refillTokens = refillTokens;
        this.refillInterval = refillInterval;
        this.availableTokens = maxTokens; // 初始可用令牌数为最大令牌数
        this.lastRefillTime = System.currentTimeMillis();
    }

    // 请求获取令牌
    public synchronized boolean tryAcquire() {
        refill(); // 尝试填充令牌

        if (availableTokens > 0) {
            availableTokens--; // 消耗一个令牌
            return true; // 允许请求通过
        }
        return false; // 拒绝请求
    }

    // 填充令牌
    private void refill() {
        long now = System.currentTimeMillis();
        long timeSinceLastRefill = now - lastRefillTime;

        // 计算可以填充的令牌数量
        long tokensToAdd = timeSinceLastRefill / refillInterval * refillTokens;
        if (tokensToAdd > 0) {
            availableTokens = Math.min(maxTokens, availableTokens + tokensToAdd);
            lastRefillTime = now; // 更新上次填充时间
        }
    }
}
